// Copyright 2020 The XLS Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef XLS_JIT_JIT_RUNTIME_H_
#define XLS_JIT_JIT_RUNTIME_H_

#include <cstddef>
#include <cstdint>
#include <memory>

#include "absl/base/thread_annotations.h"
#include "absl/status/status.h"
#include "absl/synchronization/mutex.h"
#include "absl/types/span.h"
#include "llvm/include/llvm/IR/DataLayout.h"
#include "xls/ir/type.h"
#include "xls/ir/value.h"
#include "xls/jit/llvm_type_converter.h"

namespace xls {

// JitRuntime contains routines necessary for executing code generated by the
// IR JIT. For type resolution, the JIT packs input data into and pulls
// data out of a flat character buffer, thus these routines are necessary.
class JitRuntime {
 public:
  explicit JitRuntime(llvm::DataLayout data_layout);

  // Packs the specified values into a flat buffer with the data layout
  // expected by LLVM.
  // "arg_buffers" must contain an entry corresponding to each element in
  // "args", with a matching amount of space allocated.
  absl::Status PackArgs(absl::Span<const Value> args,
                        absl::Span<Type* const> arg_types,
                        absl::Span<uint8_t* const> arg_buffers);

  // Returns a Value constructed from the data inside "buffer" whose
  // contents are laid out according to the LLVM interpretation of the passed-in
  // type.
  Value UnpackBuffer(const uint8_t* buffer, const Type* result_type);

  // Splats the value into the buffer according to the data layout expected by
  // LLVM.
  void BlitValueToBuffer(const Value& value, const Type* type,
                         absl::Span<uint8_t> buffer);

  const llvm::DataLayout& data_layout() { return data_layout_; }

  // Returns the number of bytes that should be allocated for a native LLVM
  // value storing `size` bytes with `alignment` alignment.
  //
  // Note: A user cannot just allocate `size` bytes in every scenario, as there
  //       may be alignment constraints. This method tells us how much to
  //       overallocate.
  size_t ShouldAllocateForAlignment(size_t size, int64_t alignment) const {
    return size + alignment - 1;
  }

  // Returns the number of bytes that should be allocated for a native LLVM
  // stack storing `size` bytes.
  //
  // Note: A user cannot just allocate `size` bytes in every scenario, as there
  //       may be stack alignment constraints. This method tells us how much to
  //       overallocate.
  size_t ShouldAllocateForStack(size_t size) {
    return ShouldAllocateForAlignment(
        size, data_layout_.getStackAlignment().valueOrOne().value());
  }

  // Converts the provided buffer into a native LLVM buffer by aligning to the
  // given alignment.
  //
  // May reduce the size of the buffer; if the buffer was allocated to hold at
  // least `ShouldAllocateForAlignment(size, alignment)` bytes, then the result
  // is guaranteed to hold at least `size` bytes.
  absl::Span<uint8_t> AsAligned(absl::Span<uint8_t> buffer,
                                int64_t alignment) const;

  // Converts the provided buffer into a native LLVM stack by aligning to the
  // memory model's requirements.
  //
  // May reduce the size of the buffer; if the buffer was allocated to hold at
  // least `ShouldAllocateForStack(size)` bytes, then the result is guaranteed
  // to hold at least `size` bytes.
  absl::Span<uint8_t> AsStack(absl::Span<uint8_t> buffer) {
    return AsAligned(buffer,
                     data_layout_.getStackAlignment().valueOrOne().value());
  }

  int64_t GetTypeByteSize(Type* xls_type) {
    absl::MutexLock lock(&mutex_);
    return type_converter_->GetTypeByteSize(xls_type);
  }

  int64_t GetTypeAlignment(Type* xls_type) {
    absl::MutexLock lock(&mutex_);
    return type_converter_->GetTypePreferredAlignment(xls_type);
  }

  const llvm::DataLayout& data_layout() const { return data_layout_; }

 private:
  Value UnpackBufferInternal(const uint8_t* buffer, const Type* result_type)
      ABSL_SHARED_LOCKS_REQUIRED(mutex_);
  void BlitValueToBufferInternal(const Value& value, const Type* type,
                                 absl::Span<uint8_t> buffer)
      ABSL_SHARED_LOCKS_REQUIRED(mutex_);

  mutable absl::Mutex mutex_;

  const llvm::DataLayout data_layout_;
  std::unique_ptr<llvm::LLVMContext> context_ ABSL_GUARDED_BY(mutex_);
  std::unique_ptr<LlvmTypeConverter> type_converter_ ABSL_GUARDED_BY(mutex_);
};

}  // namespace xls

#endif  // XLS_JIT_JIT_RUNTIME_H_
